import numpy as np
import heapq

# Hamiltonian循环生成器
class HamiltonianCycle:
    def __init__(self, width, height):
        self.width = width
        self.height = height
        self.cycle = self.generate_cycle()
        self.position_to_index = {pos: idx for idx, pos in enumerate(self.cycle)}
    
    def generate_cycle(self):
        """
        生成一个Hamiltonian循环路径
        使用简单的蛇形模式生成哈密顿循环
        """
        cycle = []
        # 从左上角开始
        for y in range(self.height):
            row = []
            for x in range(self.width):
                if y % 2 == 0:  # 偶数行从左到右
                    row.append((x, y))
                else:  # 奇数行从右到左
                    row.append((self.width - 1 - x, y))
            cycle.extend(row)
        
        # 添加从最后一行回到第一行的连接
        # 这里假设宽度和高度都是偶数，如果不是，需要更复杂的逻辑
        if self.height % 2 == 1:  # 如果高度是奇数
            # 最后一行是从右到左的，最后一个点是(0, height-1)
            # 需要添加一条路径回到(0, 0)
            # 这里简化处理，直接连接
            pass
        
        return cycle
    
    def get_next_position(self, current_position):
        """
        获取Hamiltonian循环中的下一个位置
        """
        if current_position not in self.position_to_index:
            # 如果当前位置不在循环中，找到最近的点
            return self.cycle[0]  # 简化处理，返回循环起点
        
        current_index = self.position_to_index[current_position]
        next_index = (current_index + 1) % len(self.cycle)
        return self.cycle[next_index]
    
    def get_cycle_distance(self, pos1, pos2):
        """
        计算在Hamiltonian循环中从pos1到pos2的距离
        """
        if pos1 not in self.position_to_index or pos2 not in self.position_to_index:
            return float('inf')  # 如果位置不在循环中，返回无穷大
        
        idx1 = self.position_to_index[pos1]
        idx2 = self.position_to_index[pos2]
        
        # 计算顺时针距离
        clockwise_dist = (idx2 - idx1) % len(self.cycle)
        # 计算逆时针距离
        counter_clockwise_dist = (idx1 - idx2) % len(self.cycle)
        
        return min(clockwise_dist, counter_clockwise_dist)
    
    def is_safe_move(self, current_position, next_position, snake_positions):
        """
        判断从当前位置移动到下一个位置是否安全
        """
        # 如果下一个位置是蛇身，则不安全
        if next_position in snake_positions[:-1]:  # 排除蛇尾，因为蛇移动时蛇尾会离开
            return False
        
        # 如果下一个位置在Hamiltonian循环中的索引比当前位置大，则安全
        # 这确保了蛇始终沿着Hamiltonian循环移动
        if current_position in self.position_to_index and next_position in self.position_to_index:
            current_index = self.position_to_index[current_position]
            next_index = self.position_to_index[next_position]
            
            # 检查是否是循环中的下一个位置或者是一个安全的捷径
            if (next_index - current_index) % len(self.cycle) == 1:
                return True
            
            # 检查是否是安全的捷径
            # 如果蛇长度小于循环中从next_position到current_position的距离，则安全
            snake_length = len(snake_positions)
            cycle_distance = (current_index - next_index) % len(self.cycle)
            if cycle_distance > snake_length:
                return True
        
        return False

# 最短路径求解器
class ShortestPathSolver:
    def __init__(self, width, height, wall_pass=False):
        self.width = width
        self.height = height
        self.wall_pass = wall_pass
    
    def get_neighbors(self, position, snake_positions):
        """
        获取给定位置的邻居位置（上、右、下、左）
        排除蛇身占据的位置
        """
        x, y = position
        neighbors = []
        
        # 上、右、下、左四个方向
        directions = [(0, -1), (1, 0), (0, 1), (-1, 0)]
        
        for dx, dy in directions:
            nx, ny = x + dx, y + dy
            
            # 处理穿墙模式
            if self.wall_pass:
                nx = nx % self.width
                ny = ny % self.height
            else:
                # 检查是否超出边界
                if nx < 0 or nx >= self.width or ny < 0 or ny >= self.height:
                    continue
            
            # 检查是否是蛇身
            if (nx, ny) in snake_positions[:-1]:  # 排除蛇尾
                continue
            
            neighbors.append((nx, ny))
        
        return neighbors
    
    def find_shortest_path(self, start, target, snake_positions):
        """
        使用A*算法找到从start到target的最短路径
        """
        # 如果起点和终点相同，直接返回
        if start == target:
            return [start]
        
        # 初始化开放列表和关闭列表
        open_list = []
        closed_set = set()
        
        # 使用曼哈顿距离作为启发函数
        def heuristic(pos):
            x1, y1 = pos
            x2, y2 = target
            
            # 处理穿墙模式
            if self.wall_pass:
                dx = min(abs(x2 - x1), self.width - abs(x2 - x1))
                dy = min(abs(y2 - y1), self.height - abs(y2 - y1))
            else:
                dx = abs(x2 - x1)
                dy = abs(y2 - y1)
            
            return dx + dy
        
        # 初始化起点
        g_score = {start: 0}
        f_score = {start: heuristic(start)}
        came_from = {}
        
        # 将起点加入开放列表
        heapq.heappush(open_list, (f_score[start], start))
        
        while open_list:
            # 获取f值最小的节点
            _, current = heapq.heappop(open_list)
            
            # 如果到达目标，构建路径并返回
            if current == target:
                path = [current]
                while current in came_from:
                    current = came_from[current]
                    path.append(current)
                return path[::-1]  # 反转路径，从起点到终点
            
            # 将当前节点加入关闭列表
            closed_set.add(current)
            
            # 检查所有邻居
            for neighbor in self.get_neighbors(current, snake_positions):
                # 如果邻居在关闭列表中，跳过
                if neighbor in closed_set:
                    continue
                
                # 计算从起点经过当前节点到达邻居的代价
                tentative_g_score = g_score[current] + 1
                
                # 如果邻居不在开放列表中，或者找到了更好的路径
                if neighbor not in g_score or tentative_g_score < g_score[neighbor]:
                    # 更新路径信息
                    came_from[neighbor] = current
                    g_score[neighbor] = tentative_g_score
                    f_score[neighbor] = g_score[neighbor] + heuristic(neighbor)
                    
                    # 将邻居加入开放列表
                    if neighbor not in [item[1] for item in open_list]:
                        heapq.heappush(open_list, (f_score[neighbor], neighbor))
        
        # 如果没有找到路径，返回None
        return None
    
    def get_next_move(self, head_position, food_position, snake_positions):
        """
        获取从蛇头到食物的下一步移动
        """
        path = self.find_shortest_path(head_position, food_position, snake_positions)
        
        if not path or len(path) < 2:
            return None  # 没有找到路径或路径太短
        
        # 返回下一步位置
        next_position = path[1]
        
        # 计算方向向量
        head_x, head_y = head_position
        next_x, next_y = next_position
        
        # 处理穿墙情况
        if self.wall_pass:
            dx = (next_x - head_x) % self.width
            if dx > self.width // 2:
                dx = dx - self.width
            
            dy = (next_y - head_y) % self.height
            if dy > self.height // 2:
                dy = dy - self.height
        else:
            dx = next_x - head_x
            dy = next_y - head_y
        
        # 标准化为单位向量
        if dx != 0 and dy != 0:
            # 如果是对角线移动，选择水平或垂直方向
            if abs(dx) > abs(dy):
                dy = 0
                dx = 1 if dx > 0 else -1
            else:
                dx = 0
                dy = 1 if dy > 0 else -1
        elif dx != 0:
            dx = 1 if dx > 0 else -1
        elif dy != 0:
            dy = 1 if dy > 0 else -1
        
        return (dx, dy)

# 混合决策系统
class HybridDecisionSystem:
    def __init__(self, width, height, wall_pass=False):
        self.hamiltonian = HamiltonianCycle(width, height)
        self.shortest_path = ShortestPathSolver(width, height, wall_pass)
        self.width = width
        self.height = height
        self.wall_pass = wall_pass
    
    def get_state_info(self, snake, food):
        """
        获取状态信息，包括Hamiltonian循环和最短路径信息
        """
        head_position = snake.get_head_position()
        food_position = food.position
        
        # 获取Hamiltonian循环中的下一个位置
        next_hamiltonian_position = self.hamiltonian.get_next_position(head_position)
        
        # 计算在Hamiltonian循环中从蛇头到食物的距离
        hamiltonian_distance = self.hamiltonian.get_cycle_distance(head_position, food_position)
        
        # 获取最短路径
        shortest_path = self.shortest_path.find_shortest_path(head_position, food_position, snake.positions)
        shortest_path_length = len(shortest_path) - 1 if shortest_path else float('inf')
        
        # 计算使用捷径的潜在收益
        shortcut_benefit = hamiltonian_distance - shortest_path_length if shortest_path else 0
        
        # 判断最短路径的下一步是否安全
        next_shortest_path_position = shortest_path[1] if shortest_path and len(shortest_path) > 1 else None
        is_shortcut_safe = False
        if next_shortest_path_position:
            is_shortcut_safe = self.hamiltonian.is_safe_move(head_position, next_shortest_path_position, snake.positions)
        
        # 获取四个方向的安全性信息
        directions = [(0, -1), (1, 0), (0, 1), (-1, 0)]  # 上、右、下、左
        direction_safety = []
        
        for dx, dy in directions:
            nx, ny = head_position[0] + dx, head_position[1] + dy
            
            # 处理穿墙模式
            if self.wall_pass:
                nx = nx % self.width
                ny = ny % self.height
            
            # 检查是否是安全移动
            is_safe = (nx, ny) not in snake.positions[:-1]  # 排除蛇尾
            if not self.wall_pass:
                is_safe = is_safe and 0 <= nx < self.width and 0 <= ny < self.height
            
            # 检查是否在Hamiltonian循环上
            is_on_hamiltonian = (nx, ny) == next_hamiltonian_position
            
            # 检查是否是最短路径的下一步
            is_on_shortest_path = next_shortest_path_position and (nx, ny) == next_shortest_path_position
            
            direction_safety.append((is_safe, is_on_hamiltonian, is_on_shortest_path))
        
        return {
            'hamiltonian_next': next_hamiltonian_position,
            'hamiltonian_distance': hamiltonian_distance,
            'shortest_path_length': shortest_path_length,
            'shortcut_benefit': shortcut_benefit,
            'is_shortcut_safe': is_shortcut_safe,
            'direction_safety': direction_safety
        }
    
    def decide_action(self, snake, food, dqn_action=None):
        """
        结合DQN、Hamiltonian循环和最短路径做出决策
        如果DQN的决策安全，则使用DQN的决策
        否则使用Hamiltonian循环作为安全策略
        """
        head_position = snake.get_head_position()
        
        # 获取状态信息
        state_info = self.get_state_info(snake, food)
        
        # 将DQN动作转换为方向
        dqn_direction = None
        if dqn_action is not None:
            if dqn_action == 0:  # 上
                dqn_direction = (0, -1)
            elif dqn_action == 1:  # 右
                dqn_direction = (1, 0)
            elif dqn_action == 2:  # 下
                dqn_direction = (0, 1)
            else:  # 左
                dqn_direction = (-1, 0)
        
        # 检查DQN的决策是否安全
        if dqn_direction:
            nx, ny = head_position[0] + dqn_direction[0], head_position[1] + dqn_direction[1]
            
            # 处理穿墙模式
            if self.wall_pass:
                nx = nx % self.width
                ny = ny % self.height
            
            # 检查是否是安全移动
            is_safe = (nx, ny) not in snake.positions[:-1]  # 排除蛇尾
            if not self.wall_pass:
                is_safe = is_safe and 0 <= nx < self.width and 0 <= ny < self.height
            
            if is_safe:
                print("使用DQN的决策_1")
                
                return dqn_direction
                # 如果DQN的决策安全，并且捷径收益大，则使用DQN的决策
                if state_info['shortcut_benefit'] > 5 and state_info['is_shortcut_safe']:
                    print("使用DQN的决策_2")
                    return dqn_direction
        print("不使用DQN的决策")
        # 如果DQN的决策不安全或不存在，或者捷径收益不大，使用Hamiltonian循环
        next_hamiltonian_position = state_info['hamiltonian_next']
        dx = next_hamiltonian_position[0] - head_position[0]
        dy = next_hamiltonian_position[1] - head_position[1]
        
        # 处理穿墙情况
        if self.wall_pass:
            if abs(dx) > 1:
                dx = -1 if dx > 0 else 1
            if abs(dy) > 1:
                dy = -1 if dy > 0 else 1
        
        return (dx, dy)